﻿#pragma once

#include <xy/render/renderer.h>

namespace xy
{

namespace render
{

// 标准着色器
template <typename shader_impl, typename primitive, std::size_t in, std::size_t out> struct shader
{
    using shader_t = shader_impl;
    using primitive_t = primitive;
    using varying_t = typename shader_t::input_varyings;
    using built_in_t = typename shader_t::built_in_varyings;
    using uniform_t = typename shader_t::uniform_varyings;
    using auxiliary_t = auxiliary_varyings_t<built_in_t, in, out>;

    // 着色器
    std::function<built_in_t(const uniform_t &, const varying_t &)> vertexShader{shader_t::vertex_shader};
    std::function<color_t(const uniform_t &, const built_in_t &)> fragmentShader{shader_t::fragment_shader};
    std::function<void(const uniform_t &, auxiliary_t &)> geometryShader{[](const uniform_t &u, auxiliary_t &a) {
        if constexpr (requires { shader_t::geometry_shader(u, a); })
            shader_t::geometry_shader(u, a);
    }};

    // 绘图参数
    uniform_t uniform{};
    std::vector<varying_t> varyings{};

    std::function<bool(const geom::vec4f &, const geom::vec4f &, const geom::vec4f &)> cullingFunc{
        renderer::back_face_culling};
    std::function<bool(std::uint32_t, std::uint32_t)> depthFunc{
        [](std::uint32_t writeZ, std::uint32_t readZ) { return writeZ < readZ; }};
};

struct point_shader_impl
{
    struct input_varyings
    {
        geom::vec3f position;
        geom::vec3f color;
    };

    struct uniform_varyings
    {
        int pointSize{10};
        int lineWidth{1};
        geom::mat4x4f vp{geom::mat4x4f::identity()};
    };

    struct built_in_varyings
    {
        geom::vec4f position;
        geom::vec3f color;

        static built_in_varyings blend(const built_in_varyings &lhs, const built_in_varyings &rhs, float t)
        {
            return {
                lhs.position * (1 - t) + rhs.position * t,
                lhs.color * (1 - t) + rhs.color * t,
            };
        }

        static built_in_varyings blend(const built_in_varyings &v0, const built_in_varyings &v1,
                                       const built_in_varyings &v2, float a, float b, float c)
        {
            return {
                v0.position * a + v1.position * b + v2.position * c,
                v0.color * a + v1.color * b + v2.color * c,
            };
        }
    };

    static built_in_varyings vertex_shader(const uniform_varyings &u, const input_varyings &v)
    {
        built_in_varyings ret;
        ret.position = u.vp * geom::cartesian_to_homogeneous(v.position);
        ret.color = v.color;
        return ret;
    }

    static color_t fragment_shader(const uniform_varyings &u, const built_in_varyings &v)
    {
        return color_t(v.color[0] * 255, v.color[1] * 255, v.color[2] * 255, 255);
    }
};

struct phong_shader_impl
{
    struct input_varyings
    {
        geom::vec3f position;
        geom::vec3f normal;
        geom::vec2f texcoord;
    };

    struct built_in_varyings
    {
        geom::vec4f position;
        geom::vec3f fragPos;
        geom::vec3f normal;
        geom::vec2f texcoord;

        // 不启用 mipmap
        // float sampleDist;
        geom::vec3f gradu;
        geom::vec3f gradv;

        static built_in_varyings blend(const built_in_varyings &lhs, const built_in_varyings &rhs, float t)
        {
            return {
                lhs.position * (1 - t) + rhs.position * t,
                lhs.fragPos * (1 - t) + rhs.fragPos * t,
                lhs.normal * (1 - t) + rhs.normal * t,
                lhs.texcoord * (1 - t) + rhs.texcoord * t,
            };
        }

        static built_in_varyings blend(const built_in_varyings &v0, const built_in_varyings &v1,
                                       const built_in_varyings &v2, float a, float b, float c)
        {
            return {
                v0.position * a + v1.position * b + v2.position * c,
                v0.fragPos * a + v1.fragPos * b + v2.fragPos * c,
                v0.normal * a + v1.normal * b + v2.normal * c,
                v0.texcoord * a + v1.texcoord * b + v2.texcoord * c,
            };
        }
    };

    struct uniform_varyings
    {
        enum Mode
        {
            smooth,
            wireframe,
            normal,
            ssao,
        };

        Mode mode{smooth};
        geom::vec3f viewPos{};
        geom::mat4x4f vp{geom::mat4x4f::identity()};
        geom::mat4x4f projection{geom::mat4x4f::identity()};
        geom::mat4x4f shadow{geom::mat4x4f::identity()};

        // 纹理采样器
        mipmap_t<4> diffuseMap;
        std::optional<image_t> specularMap;
        std::optional<image_t> normalMap;
        std::optional<image_t> shadowMap;

        // 点光源
        struct
        {
            geom::vec3f position;
            geom::vec3f direction;
            float diffuse;
            float specular;

            float constant{1.0f};
            float quadratic{0.032f};
            float occlusion{10000.0f}; // 环境光遮蔽因子
        } light;
    };

    static built_in_varyings vertex_shader(const uniform_varyings &u, const input_varyings &v)
    {
        built_in_varyings ret;
        ret.position = u.vp * geom::cartesian_to_homogeneous(v.position);
        ret.fragPos = v.position;
        ret.normal = v.normal;
        ret.texcoord = v.texcoord;
        return ret;
    }

    static color_t fragment_shader(const uniform_varyings &u, const built_in_varyings &v)
    {
        // 绘制环境光遮蔽
        if (u.mode == uniform_varyings::ssao)
            return renderer::gamma(color_t{0xffffffff} * renderer::ssao(u, v));

        // 绘制法向
        geom::vec3f normal = v.normal;
        if (u.mode == uniform_varyings::normal)
            return renderer::normal_mapping(normal);

        // 计算距离衰减
        float dis2 = (u.light.position - v.fragPos).square_modulus();
        float attenuation = 1.0f / (u.light.constant + u.light.quadratic * dis2);

        // 阴影映射
        if (u.mode == uniform_varyings::smooth && renderer::shadow(u, v))
            return color_t{0xff000000};

        // 法向贴图
        if (u.normalMap)
            normal = renderer::darboux_mapping(u, v);

        // 漫反射
        geom::vec3f lightDir = (u.light.position - v.fragPos).normalize();
        float diff = std::max(dot(normal, lightDir), 0.0f);
        color_t diffuse = renderer::mipmap_texture(u, v) * u.light.diffuse * diff;

        // 镜面反射（Blinn-Phong）
        geom::vec3f viewDir = (u.viewPos - v.fragPos).normalize();
        geom::vec3f middle = (lightDir + viewDir).normalize();
        float spec = std::pow(std::max(dot(middle, normal), 0.0f), 64.0f);

        color_t specular(0, 0, 0, 255);
        if (u.specularMap)
            specular = renderer::nearest_texture(u.specularMap.value(), v.texcoord) * u.light.specular * spec;
        else
            specular = color_t{0xffffffff} * u.light.specular * spec;

        return renderer::gamma((diffuse + specular) * attenuation);
    }
};

struct darboux_shader_impl
{
    using input_varyings = phong_shader_impl::input_varyings;
    using built_in_varyings = phong_shader_impl::built_in_varyings;

    struct uniform_varyings
    {
        int lineWidth{1};
        geom::mat4x4f vp{geom::mat4x4f::identity()};
    };

    static built_in_varyings vertex_shader(const uniform_varyings &u, const input_varyings &v)
    {
        built_in_varyings ret;
        ret.fragPos = v.position;
        ret.normal = v.normal;
        ret.texcoord = v.texcoord;
        return ret;
    }

    static color_t fragment_shader(const uniform_varyings &u, const built_in_varyings &v)
    {
        return color_t(v.normal[0] * 128, v.normal[1] * 128, v.normal[2] * 128, 255);
    }

    static void geometry_shader(const uniform_varyings &u, auxiliary_varyings_t<built_in_varyings, 3, 6> &a)
    {
        auto varying = built_in_varyings::blend(a.varyings[0], a.varyings[1], a.varyings[2], 1.0f / 3.0f, 1.0f / 3.0f,
                                                1.0f / 3.0f);
        auto [gradu, gradv] = renderer::tangent_mapping(a.varyings[0], a.varyings[1], a.varyings[2], varying.normal);

        // position 保存变换位置，normal 用来记录颜色
        a.varyings[0].position = u.vp * geom::cartesian_to_homogeneous(varying.fragPos);
        a.varyings[0].normal = geom::vec3f{1.0f, 0.0f, 0.0f};
        a.varyings[1].position = a.varyings[0].position + u.vp * geom::cartesian_to_homogeneous(gradu * 0.1f, 0.0f);
        a.varyings[1].normal = geom::vec3f{1.0f, 0.0f, 0.0f};

        a.varyings[2].position = a.varyings[0].position;
        a.varyings[2].normal = geom::vec3f{0.0f, 1.0f, 0.0f};
        a.varyings[3].position = a.varyings[0].position + u.vp * geom::cartesian_to_homogeneous(gradv * 0.1f, 0.0f);
        a.varyings[3].normal = geom::vec3f{0.0f, 1.0f, 0.0f};

        a.varyings[4].position = a.varyings[0].position;
        a.varyings[4].normal = geom::vec3f{0.0f, 0.0f, 1.0f};
        a.varyings[5].position =
            a.varyings[0].position + u.vp * geom::cartesian_to_homogeneous(varying.normal * 0.1f, 0.0f);
        a.varyings[5].normal = geom::vec3f{0.0f, 0.0f, 1.0f};
    }
};

using point_shader = shader<point_shader_impl, primitive_point, 1, 1>;       // 标准点着色器
using line_shader = shader<point_shader_impl, primitive_line, 2, 2>;         // 标准线着色器
using triangle_shader = shader<phong_shader_impl, primitive_triangle, 3, 3>; // 标准三角形着色器
using darboux_shader = shader<darboux_shader_impl, primitive_line, 3, 6>;    // darboux 着色器

} // namespace render

} // namespace xy
